/*
 * @Description: 闭环检测算法
 * @Author: Ren Qian
 * @Date: 2020-02-04 18:53:06
 */
#include "lidar_localization/mapping/loop_closing/loop_closing.hpp"

#include <cmath>
#include <algorithm>
#include <limits>

#include <pcl/common/transforms.h>
#include <pcl/io/pcd_io.h>

#include "lidar_localization/tools/logger.hpp"

#include "lidar_localization/global_defination/global_defination.h"
#include "lidar_localization/models/registration/ndt_registration.hpp"
#include "lidar_localization/models/cloud_filter/voxel_filter.hpp"
#include "lidar_localization/models/cloud_filter/no_filter.hpp"
#include "lidar_localization/tools/print_info.hpp"

namespace lidar_localization {
LoopClosing::LoopClosing() {
    InitWithConfig();
}

bool LoopClosing::InitWithConfig() {
    std::string config_file_path = WORK_SPACE_PATH + "/config/mapping/loop_closing.yaml";
    YAML::Node config_node = YAML::LoadFile(config_file_path);

    std::cout << "-----------------Init Loop-Closing Detection-------------------" << std::endl;
    InitParam(config_node);
    InitDataPath(config_node);

    InitFilter("map", map_filter_ptr_, config_node);
    InitFilter("scan", scan_filter_ptr_, config_node);

    InitLoopClosure(config_node);

    InitRegistration(registration_ptr_, config_node);

    return true;
}

bool LoopClosing::InitParam(const YAML::Node& config_node) {
    extend_frame_num_ = config_node["extend_frame_num"].as<int>();
    loop_step_ = config_node["loop_step"].as<int>();
    diff_num_ = config_node["diff_num"].as<int>();
    detect_area_ = config_node["detect_area"].as<float>();
    fitness_score_limit_ = config_node["fitness_score_limit"].as<float>();

    return true;
}

bool LoopClosing::InitDataPath(const YAML::Node& config_node) {
    std::string data_path = config_node["data_path"].as<std::string>();
    if (data_path == "./") {
        data_path = WORK_SPACE_PATH;
    }

    key_frames_path_ = data_path + "/slam_data/key_frames";
    scan_context_path_ = data_path + "/slam_data/scan_context";

    return true;
}

bool LoopClosing::InitFilter(
    std::string filter_user, 
    std::shared_ptr<CloudFilterInterface>& filter_ptr, 
    const YAML::Node& config_node
) {
    std::string filter_mothod = config_node[filter_user + "_filter"].as<std::string>();
    std::cout << "\tFilter Method for " << filter_user << ": " << filter_mothod << std::endl;

    if (filter_mothod == "voxel_filter") {
        filter_ptr = std::make_shared<VoxelFilter>(config_node[filter_mothod][filter_user]);
    } else if (filter_mothod == "no_filter") {
        filter_ptr =std::make_shared<NoFilter>();
    } else {
        AERROR << "Filter method " << filter_mothod << " for " << filter_user << " NOT FOUND!";
        return false;
    }

    return true;
}

bool LoopClosing::InitLoopClosure(const YAML::Node& config_node) {
    // get loop closure config:
    loop_closure_method_ = config_node["loop_closure_method"].as<std::string>();
    
    if(loop_closure_method_=="scan_context")
    {
        scan_context_manager_ptr_ = std::make_shared<ScanContextManager>(config_node[loop_closure_method_]);
    }
    else if(loop_closure_method_=="std_detector")
    {
        ConfigSetting config_setting;
        // pre-preocess
        config_setting.ds_size_ = config_node["std_detector"]["ds_size_"].as<double>();
        config_setting.maximum_corner_num_ = config_node["std_detector"]["maximum_corner_num_"].as<int>();
        // key points
        config_setting.plane_merge_normal_thre_ = config_node["std_detector"]["plane_merge_normal_thre_"].as<double>();
        config_setting.plane_detection_thre_ = config_node["std_detector"]["plane_detection_thre_"].as<double>();
        config_setting.voxel_size_ = config_node["std_detector"]["voxel_size_"].as<double>();
        config_setting.voxel_init_num_ = config_node["std_detector"]["voxel_init_num_"].as<int>();
        config_setting.proj_image_resolution_ = config_node["std_detector"]["proj_image_resolution_"].as<double>();
        config_setting.proj_dis_min_ = config_node["std_detector"]["proj_dis_min_"].as<double>();
        config_setting.proj_dis_max_ = config_node["std_detector"]["proj_dis_max_"].as<double>();
        config_setting.corner_thre_ = config_node["std_detector"]["corner_thre_"].as<double>();

        // std descriptor
        config_setting.descriptor_near_num_ = config_node["std_detector"]["descriptor_near_num_"].as<int>();
        config_setting.descriptor_min_len_ = config_node["std_detector"]["descriptor_min_len_"].as<double>();
        config_setting.descriptor_max_len_ = config_node["std_detector"]["descriptor_max_len_"].as<double>();
        config_setting.non_max_suppression_radius_ = config_node["std_detector"]["non_max_suppression_radius_"].as<double>();
        config_setting.std_side_resolution_ = config_node["std_detector"]["std_side_resolution_"].as<double>();

        // candidate search
        config_setting.skip_near_num_ = config_node["std_detector"]["skip_near_num_"].as<int>();
        config_setting.candidate_num_ = config_node["std_detector"]["candidate_num_"].as<int>();
        config_setting.descriptor_max_len_ = config_node["std_detector"]["descriptor_max_len_"].as<int>();
        config_setting.rough_dis_threshold_ = config_node["std_detector"]["rough_dis_threshold_"].as<double>();
        config_setting.vertex_diff_threshold_ = config_node["std_detector"]["vertex_diff_threshold_"].as<double>();  

        config_setting.icp_threshold_ = config_node["std_detector"]["icp_threshold_"].as<double>();
        config_setting.normal_threshold_ = config_node["std_detector"]["normal_threshold_"].as<double>();
        config_setting.dis_threshold_ = config_node["std_detector"]["dis_threshold_"].as<double>(); 

        _std_manager_ptr = std::make_shared<STDescManager>(config_setting);
        AWARN << "loop_closure_method_" << ":" << loop_closure_method_ << std::endl;
        std::cout << "Sucessfully load parameters:" << std::endl;
        std::cout << "----------------Main Parameters-------------------"
                    << std::endl;
        std::cout << "voxel size:" << config_setting.voxel_size_ << std::endl;
        std::cout << "loop detection threshold: " << config_setting.icp_threshold_
                    << std::endl;
        std::cout << "sub-frame number: " << config_setting.sub_frame_num_
                    << std::endl;
        std::cout << "candidate number: " << config_setting.candidate_num_
                    << std::endl;
        std::cout << "maximum corners size: " << config_setting.maximum_corner_num_
                    << std::endl;
    }
    // create instance:

    return true;
}

bool LoopClosing::InitRegistration(std::shared_ptr<RegistrationInterface>& registration_ptr, const YAML::Node& config_node) {
    std::string registration_method = config_node["registration_method"].as<std::string>();
    std::cout << "\tPoint Cloud Registration Method: " << registration_method << std::endl;

    if (registration_method == "NDT") {
        registration_ptr = std::make_shared<NDTRegistration>(config_node[registration_method]);
    } else {
        AERROR << "Registration method " << registration_method << " NOT FOUND!";
        return false;
    }

    return true;
}

bool LoopClosing::Update(
    const CloudData &key_scan, 
    const KeyFrame &key_frame, 
    const KeyFrame &key_gnss
) {
    static int key_frame_index = 0;
    static float yaw_change_in_rad = 0.0f;

    has_new_loop_pose_ = false;

    all_key_frames_.push_back(key_frame);
    all_key_gnss_.push_back(key_gnss);
    //设置了通过文件选择回环的方式
    if(loop_closure_method_ == "scan_context"){

        scan_context_manager_ptr_->Update(key_scan, key_gnss);
        if (!ScDetectNearestKeyFrame(key_frame_index, yaw_change_in_rad)){
            return false;
        }

        if (!CloudRegistration(key_frame_index, yaw_change_in_rad))
            return false;
    }else if(loop_closure_method_ == "std_detector"){
        
        if (!StdDetectNearestKeyFrame(key_frame_index, yaw_change_in_rad, key_scan,key_gnss)){
            return false;
        }
    }




    has_new_loop_pose_ = true;
    return true;
}

bool LoopClosing::ScDetectNearestKeyFrame(
    int& key_frame_index,
    float& yaw_change_in_rad
) {
    static int skip_cnt = 0;
    static int skip_num = loop_step_;
    
    // only perform loop closure detection for every skip_num key frames:
    if (++skip_cnt < skip_num)
        return false;

    #ifndef SCAN_CONTEXT
        // generate loop-closure proposal using scan context match:
        std::pair<int, float> proposal = scan_context_manager_ptr_->DetectLoopClosure();
        const int proposed_key_frame_id = proposal.first;
        const float proposed_yaw_change = proposal.second;

        // check proposal validity:
        if (ScanContextManager::NONE == proposed_key_frame_id) {
            return false;
        }

        // check RTK position difference:
        const KeyFrame &current_key_frame = all_key_gnss_.back();
        const KeyFrame &proposed_key_frame = all_key_gnss_.at(proposed_key_frame_id);

        Eigen::Vector3f translation = (
            current_key_frame.pose.block<3, 1>(0, 3) - proposed_key_frame.pose.block<3, 1>(0, 3)
        );
        float key_frame_distance = translation.norm();
    #else
        // total number of GNSS/IMU key frame poses:
        const size_t N = all_key_gnss_.size();

        // ensure valid loop closure match:
        if (
            N < static_cast<size_t>(diff_num_ + 1)
        )
            return false;

        const KeyFrame &current_key_frame = all_key_gnss_.back();

        int proposed_key_frame_id = ScanContextManager::NONE;
        // this orientation compensation is not available for GNSS/IMU proposal:
        const float proposed_yaw_change = 0.0f;
        float key_frame_distance = std::numeric_limits<float>::max();
        for (size_t i = 0; i < N - 1; ++i) {
            // ensure key frame seq. distance:
            if (N < static_cast<size_t>(i + diff_num_))
                break;
            
            const KeyFrame &proposed_key_frame = all_key_gnss_.at(i);

            Eigen::Vector3f translation = (
                current_key_frame.pose.block<3, 1>(0, 3) - proposed_key_frame.pose.block<3, 1>(0, 3)
            );
            float distance = translation.norm();

            // get closest proposal:
            if (distance < key_frame_distance) {
                key_frame_distance = distance;
                proposed_key_frame_id = i;
            }
        }
    #endif

    // this is needed for valid local map build:
    if (proposed_key_frame_id < extend_frame_num_)
        return false;

    // update detection interval:
    skip_cnt = 0;
    skip_num = static_cast<int>(key_frame_distance);
    if (key_frame_distance > detect_area_) {
        skip_num = std::max((int)(key_frame_distance / 2.0), loop_step_);
        return false;
    } else {
        key_frame_index = proposed_key_frame_id;
        yaw_change_in_rad = proposed_yaw_change;

        skip_num = loop_step_;
        return true;
    }
}

bool LoopClosing::StdDetectNearestKeyFrame(
      int& key_frame_index,
      float& yaw_change_in_rad,
      const CloudData &key_scan,
      const KeyFrame &key_gnss
    ) {
    static int skip_cnt = 0;
    static int skip_num = loop_step_;
    _key_cloud_Ind ++;
    pcl::PointCloud<pcl::PointXYZI>::Ptr current_cloud(
    new pcl::PointCloud<pcl::PointXYZI>());
    Eigen::Vector3d trans = key_gnss.pose.block<3,1>(0,3).cast<double>();
    Eigen::Matrix3d rota = key_gnss.pose.block<3,3>(0,0).cast<double>();
    for (std::size_t i = 0; i < key_scan.cloud_ptr->size(); i++) {
      pcl::PointXYZI point;
      point.x = key_scan.cloud_ptr->at(i).x;
      point.y = key_scan.cloud_ptr->at(i).y;
      point.z = key_scan.cloud_ptr->at(i).z;
      point.intensity = 0;
      Eigen::Vector3d pv = point2vec(point);
      Eigen::Vector3d pc = rota * pv + trans;
      point = vec2point(pc);
      current_cloud->push_back(point);
    }
    // pcl::copyPointCloud(*key_scan.cloud_ptr, *current_cloud);

    // pcl::transformPointCloud(*current_cloud, *current_cloud, 
    //                         key_gnss.pose);
    std::vector<STDesc> stds_vec;
    _std_manager_ptr->GenerateSTDescs(current_cloud, stds_vec);
    // only perform loop closure detection for every skip_num key frames:
    if (++skip_cnt < skip_num)
    {
        _std_manager_ptr->AddSTDescs(stds_vec);
        return false;
    }
    std::pair<int, double> search_result(-1, 0);
    std::pair<Eigen::Vector3d, Eigen::Matrix3d> loop_transform;
    loop_transform.first << 0, 0, 0;
    loop_transform.second = Eigen::Matrix3d::Identity();
    std::vector<std::pair<STDesc, STDesc>> loop_std_pair;
    _std_manager_ptr->SearchLoop(stds_vec, search_result, loop_transform,
                                loop_std_pair);
    _std_manager_ptr->AddSTDescs(stds_vec);
    const int proposed_key_frame_id = search_result.first;
    // Eigen::Matrix3f A = loop_transform.second.cast<float>();
    // const float proposed_yaw_change = A.eulerAngles(2,1,0)[2];
    // check proposal validity:
    //   std::cout << "Key Frame id:" << keyCloudInd
    //             << ", cloud size: " << current_cloud->size() << std::endl;
            
    if (search_result.first < 0) {
        return false;
    }

    // check RTK position difference:
    const KeyFrame &current_key_frame = all_key_gnss_.back();
    const KeyFrame &proposed_key_frame = all_key_gnss_.at(proposed_key_frame_id);

    Eigen::Vector3f translation = (
        current_key_frame.pose.block<3, 1>(0, 3) - proposed_key_frame.pose.block<3, 1>(0, 3)
    );
    float key_frame_distance = translation.norm();

    // this is needed for valid local map build:
    if (proposed_key_frame_id < extend_frame_num_)
        return false;

    // update detection interval:
    skip_cnt = 0;
    skip_num = static_cast<int>(key_frame_distance);
    if (key_frame_distance > detect_area_) {
        skip_num = std::max((int)(key_frame_distance / 2.0), loop_step_);
        return false;
    } else {
        key_frame_index = proposed_key_frame_id;
        // yaw_change_in_rad = proposed_yaw_change;
        std::cout << "[Loop Detection] triggle loop: " << _key_cloud_Ind-1 << "--"
                << search_result.first << ", score:" << search_result.second
                << std::endl;
        current_loop_pose_.pose.block<3,3>(0,0) = loop_transform.second.cast<float>();
        current_loop_pose_.pose.block<3,1>(0,3) = loop_transform.first.cast<float>();
        skip_num = loop_step_;
        return true;
    }
}

bool LoopClosing::CloudRegistration(
    const int key_frame_index,
    const float yaw_change_in_rad
) {
    // 生成地图
    CloudData::CLOUD_PTR map_cloud_ptr(new CloudData::CLOUD());
    Eigen::Matrix4f map_pose = Eigen::Matrix4f::Identity();
    JointMap(
        key_frame_index, yaw_change_in_rad, 
        map_cloud_ptr, map_pose
    );

    // 生成当前scan
    CloudData::CLOUD_PTR scan_cloud_ptr(new CloudData::CLOUD());
    Eigen::Matrix4f scan_pose = Eigen::Matrix4f::Identity();
    JointScan(scan_cloud_ptr, scan_pose);

    // 匹配
    Eigen::Matrix4f result_pose = Eigen::Matrix4f::Identity();
    Registration(map_cloud_ptr, scan_cloud_ptr, scan_pose, result_pose);

    // 计算相对位姿
    current_loop_pose_.pose = map_pose.inverse() * result_pose;

    // 判断是否有效
    if (registration_ptr_->GetFitnessScore() > fitness_score_limit_)
        return false;
    
    static int loop_close_cnt = 0;
    loop_close_cnt ++;

    AINFO << std::endl
              << "[ICP Registration] Loop-Closure Detected " 
              << current_loop_pose_.index0 << "<-->" << current_loop_pose_.index1 << std::endl 
              << "\tFitness Score " << registration_ptr_->GetFitnessScore() << std::endl 
              << std::endl;

    return true;
}

bool LoopClosing::JointMap(
    const int key_frame_index, const float yaw_change_in_rad,
    CloudData::CLOUD_PTR& map_cloud_ptr, Eigen::Matrix4f& map_pose
) {
    // init map pose as loop closure pose:
    map_pose = all_key_gnss_.at(key_frame_index).pose;

    // apply yaw change estimation from scan context match:
    Eigen::AngleAxisf orientation_change(yaw_change_in_rad, Eigen::Vector3f::UnitZ());
    map_pose.block<3, 3>(0, 0) = map_pose.block<3, 3>(0, 0) * orientation_change.toRotationMatrix();

    current_loop_pose_.index0 = all_key_frames_.at(key_frame_index).index;
    
    // create local map:
    Eigen::Matrix4f pose_to_gnss = map_pose * all_key_frames_.at(key_frame_index).pose.inverse();
    for (int i = key_frame_index - extend_frame_num_; i < key_frame_index + extend_frame_num_; ++i) {
        // a. load back surrounding key scan:
        std::string file_path = key_frames_path_ + "/key_frame_" + std::to_string(all_key_frames_.at(i).index) + ".pcd";
        CloudData::CLOUD_PTR cloud_ptr(new CloudData::CLOUD());
        pcl::io::loadPCDFile(file_path, *cloud_ptr);
        
        // b. transform surrounding key scan to map frame:
        Eigen::Matrix4f cloud_pose = pose_to_gnss * all_key_frames_.at(i).pose;
        pcl::transformPointCloud(*cloud_ptr, *cloud_ptr, cloud_pose);

        *map_cloud_ptr += *cloud_ptr;
    }
    // pre-process current map:
    map_filter_ptr_->Filter(map_cloud_ptr, map_cloud_ptr);

    return true;
}

bool LoopClosing::JointScan(CloudData::CLOUD_PTR& scan_cloud_ptr, Eigen::Matrix4f& scan_pose) {
    // set scan pose as GNSS estimation:
    scan_pose = all_key_gnss_.back().pose;
    current_loop_pose_.index1 = all_key_frames_.back().index;
    current_loop_pose_.time = all_key_frames_.back().time;

    // load back current key scan:
    std::string file_path = key_frames_path_ + "/key_frame_" + std::to_string(all_key_frames_.back().index) + ".pcd";
    std::cout << " load " << file_path << std::endl;
    pcl::io::loadPCDFile(file_path, *scan_cloud_ptr);

    // pre-process current scan:
    scan_filter_ptr_->Filter(scan_cloud_ptr, scan_cloud_ptr);

    return true;
}

bool LoopClosing::Registration(CloudData::CLOUD_PTR& map_cloud_ptr, 
                               CloudData::CLOUD_PTR& scan_cloud_ptr, 
                               Eigen::Matrix4f& scan_pose, 
                               Eigen::Matrix4f& result_pose) {
    // point cloud registration:
    CloudData::CLOUD_PTR result_cloud_ptr(new CloudData::CLOUD());
    registration_ptr_->SetInputTarget(map_cloud_ptr);
    registration_ptr_->ScanMatch(scan_cloud_ptr, scan_pose, result_cloud_ptr, result_pose);

    return true;
}

bool LoopClosing::HasNewLoopPose() {
    return has_new_loop_pose_;
}

LoopPose& LoopClosing::GetCurrentLoopPose() {
    return current_loop_pose_;
}

bool LoopClosing::Save(void) {
    return scan_context_manager_ptr_->Save(scan_context_path_);
}

}